# -*- coding: utf-8 -*-
from stateSecret.sm2_check import CryptSM2
from base64 import b64encode, b64decode
from stateSecret.sm2utils import priKey
import random

# sm2的公私钥

sm2_crypt = CryptSM2()


class sm2Encrypt:
    @staticmethod
    def make_sm2() -> dict:
        """
        生成sm2公私钥
        :return: sm2公私钥
        """
        pubKey = priKey.publicKey()
        return {"prvKey": priKey.toString(), "pubKey": pubKey.toString(compressed=False)}

    @staticmethod
    def encrypt(info: str, pubKey: str) -> str:
        """
        sm2公钥加密
        :param info: 原始数据
        :param pubKey: 公钥
        :return: 加密结果
        """
        encode_info = sm2_crypt.encrypt(info.encode(encoding="utf-8"), pubKey)
        encode_info = b64encode(encode_info).decode()  # 将二进制bytes通过base64编码
        return encode_info

    @staticmethod
    def decrypt(info: str, prvKey: str) -> str:
        """
        私钥解密
        :param info: 公钥加密后的结果
        :param prvKey: 私钥
        :return: 私钥解密后的结果
        """
        decode_info = b64decode(info.encode())  # 通过base64解码成二进制bytes
        decode_info = sm2_crypt.decrypt(decode_info, prvKey).decode(encoding="utf-8")
        return decode_info

    @staticmethod
    def sign(info: str, prvKey: str) -> str:
        """
        sm2私钥签名
        :param info: 原始数据
        :param prvKey: 私钥
        :return: 签名后的结果
        """
        decode_info = info.encode(encoding="utf-8")
        number = str(random.randint(0, 100))
        number_info = number.encode(encoding="utf-8")
        decode_info = sm2_crypt.sign(decode_info, number_info, prvKey)
        return decode_info

    @staticmethod
    def verify(info: str, origin_data: str, pubKey: str) -> bool:
        """
        sm2验签
        :param info: sm2签名的结果
        :param origin_data: 原始数据
        :param pubKey: 公钥
        :return: true 验签成功 false 验签失败
        """
        decode_info = info.encode()
        origin_data_info = origin_data.encode(encoding="utf-8")
        decode_info = sm2_crypt.verify(decode_info, origin_data_info, pubKey)
        return decode_info

sm2 = sm2Encrypt()
if __name__ == "__main__":
    sm2_key = sm2.make_sm2()
    SM2_PRIVATE_KEY, SM2_PUBLIC_KEY = sm2_key["prvKey"], sm2_key["pubKey"]
    print("sm2公钥---------", SM2_PUBLIC_KEY)
    print("sm2私钥---------", SM2_PRIVATE_KEY)
    origin_pwd = '123456'
    sm2 = sm2Encrypt()
    encrypt_pwd = sm2.encrypt(origin_pwd, SM2_PUBLIC_KEY)
    print("sm2公钥加密---------", encrypt_pwd)
    decrypt_pwd = sm2.decrypt(encrypt_pwd, SM2_PRIVATE_KEY)
    print("sm2私钥解密---------", decrypt_pwd)
    sign = sm2.sign(origin_pwd, SM2_PRIVATE_KEY)
    print("sm2私钥签名---------", sign)
    print("sm2公钥验签---------", sm2.verify(sign, origin_pwd, SM2_PUBLIC_KEY))
